#!/usr/bin/env python3
"""
DataSaver文件夹输出功能演示
展示split_data的文件夹输出支持
"""

import sys
import os
import tempfile
import json

# 添加项目路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from plugins.data_analysis.data_analysis import DataSplitter, DataSaver

def create_sample_data():
    """创建示例数据"""
    return [
        {"name": "张三", "department": "销售部", "age": 28, "salary": 8000, "performance": 85},
        {"name": "李四", "department": "技术部", "age": 32, "salary": 12000, "performance": 92},
        {"name": "王五", "department": "销售部", "age": 29, "salary": 8500, "performance": 88},
        {"name": "赵六", "department": "人事部", "age": 26, "salary": 7000, "performance": 78},
        {"name": "钱七", "department": "技术部", "age": 35, "salary": 15000, "performance": 95},
        {"name": "孙八", "department": "销售部", "age": 31, "salary": 9000, "performance": 82},
        {"name": "周九", "department": "人事部", "age": 30, "salary": 7500, "performance": 80},
        {"name": "吴十", "department": "技术部", "age": 33, "salary": 13000, "performance": 90},
        {"name": "郑十一", "department": "财务部", "age": 34, "salary": 11000, "performance": 87},
        {"name": "王十二", "department": "财务部", "age": 27, "salary": 9500, "performance": 83}
    ]

def demo_folder_output():
    """演示文件夹输出功能"""
    print("=== DataSaver文件夹输出功能演示 ===\n")
    
    # 创建示例数据
    sample_data = create_sample_data()
    print(f"📊 原始数据: {len(sample_data)} 条记录")
    
    with tempfile.TemporaryDirectory() as temp_dir:
        print(f"📁 临时输出目录: {temp_dir}\n")
        
        # 1. 按部门拆分的文件夹输出（JSON格式）
        print("1️⃣ 按部门拆分 + JSON格式文件夹输出")
        dept_output_dir = os.path.join(temp_dir, "departments_json")
        
        # 拆分数据
        dept_splitter = DataSplitter("dept_splitter", {
            "config": {
                "split_strategy": "field_value",
                "split_field": "department"
            }
        })
        dept_split_result = dept_splitter.execute(sample_data)
        
        # 保存到文件夹
        dept_saver = DataSaver("dept_saver", {
            "config": {
                "format": "json",
                "output_path": dept_output_dir,
                "split_output_folder": True,
                "file_prefix": "dept",
                "include_metadata": True,
                "pretty_print": True
            }
        })
        dept_save_result = dept_saver.execute(dept_split_result)
        
        if dept_save_result.get("saved"):
            print(f"   ✅ 成功保存到: {dept_save_result['file_path']}")
            saved_files = dept_save_result.get('metadata', {}).get('saved_files', [])
            print(f"   📄 生成文件: {len(saved_files)} 个")
            for file_info in saved_files:
                print(f"      - {os.path.basename(file_info['file_path'])} ({file_info['record_count']} 条记录)")
        print()
        
        # 2. 按比例拆分的文件夹输出（CSV格式）
        print("2️⃣ 按比例拆分 + CSV格式文件夹输出")
        ratio_output_dir = os.path.join(temp_dir, "train_test_csv")
        
        # 拆分数据（训练集70%，测试集30%）
        ratio_splitter = DataSplitter("ratio_splitter", {
            "config": {
                "split_strategy": "ratio",
                "ratios": [0.7, 0.3],
                "group_names": ["train", "test"],
                "shuffle": True,
                "random_seed": 42
            }
        })
        ratio_split_result = ratio_splitter.execute(sample_data)
        
        # 保存到文件夹
        ratio_saver = DataSaver("ratio_saver", {
            "config": {
                "format": "csv",
                "output_path": ratio_output_dir,
                "split_output_folder": True,
                "file_prefix": "dataset",
                "include_metadata": False
            }
        })
        ratio_save_result = ratio_saver.execute(ratio_split_result)
        
        if ratio_save_result.get("saved"):
            print(f"   ✅ 成功保存到: {ratio_save_result['file_path']}")
            saved_files = ratio_save_result.get('metadata', {}).get('saved_files', [])
            print(f"   📄 生成文件: {len(saved_files)} 个")
            for file_info in saved_files:
                print(f"      - {os.path.basename(file_info['file_path'])} ({file_info['record_count']} 条记录)")
        print()
        
        # 3. 按条件拆分的文件夹输出（TXT格式）
        print("3️⃣ 按条件拆分 + TXT格式文件夹输出")
        condition_output_dir = os.path.join(temp_dir, "performance_groups_txt")
        
        # 拆分数据（按绩效分组）
        condition_splitter = DataSplitter("perf_splitter", {
            "config": {
                "split_strategy": "condition",
                "condition_field": "performance",
                "conditions": {
                    "high": "performance >= 90",
                    "medium": "performance >= 80 and performance < 90",
                    "low": "performance < 80"
                }
            }
        })
        condition_split_result = condition_splitter.execute(sample_data)
        
        # 保存到文件夹
        condition_saver = DataSaver("condition_saver", {
            "config": {
                "format": "txt",
                "output_path": condition_output_dir,
                "split_output_folder": True,
                "file_prefix": "performance",
                "pretty_print": True
            }
        })
        condition_save_result = condition_saver.execute(condition_split_result)
        
        if condition_save_result.get("saved"):
            print(f"   ✅ 成功保存到: {condition_save_result['file_path']}")
            saved_files = condition_save_result.get('metadata', {}).get('saved_files', [])
            print(f"   📄 生成文件: {len(saved_files)} 个")
            for file_info in saved_files:
                print(f"      - {os.path.basename(file_info['file_path'])} ({file_info['record_count']} 条记录)")
        print()
        
        # 4. 对比：不使用文件夹输出
        print("4️⃣ 对比：传统单文件输出")
        single_file = os.path.join(temp_dir, "all_data.json")
        
        # 使用相同的拆分结果
        single_saver = DataSaver("single_saver", {
            "config": {
                "format": "json",
                "output_path": single_file,
                "split_output_folder": False,  # 禁用文件夹输出
                "include_metadata": True,
                "pretty_print": True
            }
        })
        single_save_result = single_saver.execute(dept_split_result)
        
        if single_save_result.get("saved"):
            print(f"   ✅ 成功保存到: {single_save_result['file_path']}")
            print(f"   📊 总记录数: {single_save_result['record_count']}")
            print(f"   📋 数据类型: 单个文件包含所有分组")
        print()
        
        # 总结
        print("📊 功能总结:")
        print("   ✅ split_output_folder=True: 每个分组保存为独立文件")
        print("   ✅ 支持多种格式: JSON、CSV、TXT")
        print("   ✅ 自动文件命名: 支持file_prefix前缀")
        print("   ✅ 兼容传统模式: split_output_folder=False 保持原有行为")
        print("   ✅ 详细统计信息: 记录数、分组信息、文件列表")
        print()
        
        print("💡 使用场景:")
        print("   🎯 机器学习: 训练集/测试集分别保存")
        print("   🎯 数据归档: 按类别分别存储")
        print("   🎯 报表生成: 不同维度数据独立文件")
        print("   🎯 数据备份: 结构化分组备份")

if __name__ == "__main__":
    demo_folder_output()